#!/usr/bin/env python
# encoding: utf-8
"""
primer3.py

Created by Brant Faircloth on 2007-01-22.
Copyright (c) 2007 Brant C. Faircloth. All rights reserved.

Released under the GNU General Public License 2.0 
http://www.gnu.org/licenses/gpl.html

"""

import sys, os, csv, re, pdb
from Bio import Fasta
#----------------------------------------------------------------------
# get files with repeats from excel spreadsheet output by msatCommander
#----------------------------------------------------------------------
def getData(dataFile):
    """get the freaking data"""
    dataDialect = csv.Sniffer().sniff(dataFile)                                 # try to determine dialect automagically
    headerTrue = csv.Sniffer().has_header(dataFile)                             # try to determine if file has header row
    msatData = open(dataFile,'r')                                               # open datafile for read
    msatReader = csv.reader(msatData, delimiter=',', dialect = dataDialect)     # setup csv.reader
    if headerTrue:
        msatReader.next()
    msatSearch = {}
    for row in msatReader:
        clone,desc,repeat,loc,start,bs,end = row
        msatSearch[clone]=[repeat, int(start), int(end)]
    msatData.close()
    return msatSearch

class primerRun:
    """this class will run the functions associated with primer3_core"""
    
    def __init__(self, s): 
        """create DNA instance initialized to string s.""" 
        self.seq = s
        reg = re.compile('[^ACGTN]+')           # find non-ACGT bases
        self.seq = reg.sub('N',self.seq)        # convert non-ACGT bases to N
    
    def offset(self, start, length, offset):
        """function adds length to target region"""
        if start - offset <= 0: pass
        else: start -= offset                   # subtract 10 bases from target start (for non-specific crap)
        if length + offset > len(self.seq): pass
        else: length += offset*2                # add 20 bases to target length (for non-specific crap)
        self.start = start
        self.length = length                    
    
    def format(self, start, length):
        """function builds temporary output file to be used by primer 3, subclassed below"""
        self.offset(start, length, 10)
        outFile=open('primer3_temp.txt','w')
        outFile.write(('SEQUENCE=%s\n') % (self.seq))
        outFile.write(('TARGET=%s,%s\n') % (self.start, self.length))
        features = ["PRIMER_PRODUCT_SIZE_RANGE=150-450",
            "PRIMER_MIN_TM=57.0",
            "PRIMER_MAX_TM=62.0",
            "PRIMER_OPT_TM=60.0",
            "PRIMER_MAX_DIFF_TM=5.0",
            "PRIMER_TM_SANTALUCIA=1",
            "PRIMER_SALT_CONC=50.0",
            "PRIMER_DIVALENT_CONC=1.5",
            "PRIMER_SALT_CORRECTIONS=1",
            "PRIMER_DNTP_CONC=0.125",
            "PRIMER_OPT_SIZE=19",
            "PRIMER_MIN_SIZE=17",
            "PRIMER_MIN_GC=35.0",
            "PRIMER_OPT_GC_PERCENT=50.0",
            "PRIMER_MAX_GC=75.0",
            "PRIMER_GC_CLAMP=1",
            "PRIMER_SELF_ANY=6.0",
            "PRIMER_SELF_END=2.5",
            "PRIMER_MAX_END_STABILITY=8.0",
            "PRIMER_MAX_POLY_X=4",
            "PRIMER_EXPLAIN_FLAG=1",
            "PRIMER_NUM_RETURN=4",
            "="
            ]
        for feat in features:
            outFile.write(('%s\n') % (feat))
        outFile.close()
        #------------------------
        # End primer3 file values
        #------------------------
    
    def search(self, clone, primerWriter):
        pathToPrettyPrimer3 = (('~/bin/primer3_core -format_output < primer3_temp.txt > %s') % (os.getcwd() + '/primer3/' + clone +'_primer.txt'))
        #pathToPrettyPrimer3 = (('../src/primer3_core -format_output < ../scripts/temp.txt > %s') % (clone +'.txt'))
        os.popen(pathToPrettyPrimer3)             # this should output pretty format
        #pathToPrimer3 = '../src/primer3_core < ../scripts/temp.txt'
        pathToPrimer3 = '~/bin/primer3_core < primer3_temp.txt'
        #--------------------------------------------------------------------------------------
        # here we call primer3, read the stdout from primer3, and put that output into a
        # string.  The chop up the string and put it in a dictionary - ahhhhh, dictionaries.
        #--------------------------------------------------------------------------------------        
        primer3Output=os.popen(pathToPrimer3).read().split('\n')
        primer0specs = {
        'PRIMER_RIGHT_SEQUENCE':None,
        'PRIMER_RIGHT_TM':None,
        'PRIMER_RIGHT_GC_PERCENT':None,
        'PRIMER_LEFT_SEQUENCE':None,
        'PRIMER_LEFT_TM':None,
        'PRIMER_LEFT_GC_PERCENT':None,
        'PRIMER_RIGHT_SELF_ANY':None,
        'PRIMER_RIGHT_SELF_END':None,
        'PRIMER_RIGHT_END_STABILITY':None,
        'PRIMER_RIGHT_PENALTY':None,
        'PRIMER_RIGHT_EXPLAIN':None,
        'PRIMER_LEFT_SELF_ANY':None,
        'PRIMER_LEFT_SELF_END':None,
        'PRIMER_LEFT_END_STABILITY':None,
        'PRIMER_LEFT_PENALTY':None,	
        'PRIMER_LEFT_EXPLAIN':None
        }
        primer1specs,primer2specs,primer3specs,primer4specs = primer0specs.copy(),primer0specs.copy(),primer0specs.copy(),primer0specs.copy()
        listOfDicts = [primer0specs,primer1specs,primer2specs,primer3specs,primer4specs]
        compileStatements = [
        '[A-Z]+_(RIGHT|LEFT)_[A-Z_]*=',
        '[A-Z]+_(RIGHT|LEFT)_1_[A-Z_]*=',
        '[A-Z]+_(RIGHT|LEFT)_2_[A-Z_]*=',
        '[A-Z]+_(RIGHT|LEFT)_3_[A-Z_]*=',
        '[A-Z]+_(RIGHT|LEFT)_4_[A-Z_]*='
        ]
        i=0
        for primerDict in listOfDicts:
            reg = re.compile(compileStatements[i])
            #print 'primer 3 output', primer3Output
            for item in primer3Output:
                m = reg.match(item)
                if m:
                    #m.group()
                    splitItem = item.split('=')
                    primerDict[splitItem[0]]=splitItem[1]
                else: pass
            i+=1
        outputdata = [
            clone,
            primer0specs['PRIMER_LEFT_SEQUENCE'],
            primer0specs['PRIMER_LEFT_TM'],
            primer0specs['PRIMER_LEFT_GC_PERCENT'],
            primer0specs['PRIMER_RIGHT_SEQUENCE'],
            primer0specs['PRIMER_RIGHT_TM'],
            primer0specs['PRIMER_RIGHT_GC_PERCENT'],
            ]
        print 'output data is ', outputdata
        if outputdata[1]:
            primerWriter.writerow(outputdata)
        
        

def main():
    msatSearchResults = 'cleanMsatContig.csv'
    fastaFilename = 'output.fsa'
    try:
        os.mkdir(os.getcwd()+'/primer3')
    except:
        pass
    primerDataFile = 'primer3/primerNoTag.csv'
    searchResults = getData(msatSearchResults) 
    infile = open(fastaFilename,'r')
    parser = Fasta.RecordParser()
    iterator = Fasta.Iterator(infile, parser)
    seqDict = {}
    while 1:
        #pdb.set_trace()
        seq = iterator.next()
        if not seq:
            break
            infile.close()
        #------------------------------------------------------------------------
        # this needs to be removed for other plates since the contig names differ
        #------------------------------------------------------------------------
        else:
            if seq.title in searchResults.keys():
                searchResults[seq.title].append(seq.sequence)       #append sequence to last element of each list for dictionary (searchResults)
    
    primerData = open(primerDataFile,'w')                               # open datafile for read
    primerWriter = csv.writer(primerData, dialect = 'excel')
    header = [
        'CLONE',
        'PRIMER_LEFT',
        'LEFT_TM',
        'LEFT_GC_PERCENT',
        'PRIMER_RIGHT',
        'RIGHT_TM',
        'RIGHT_GC_PERCENT',
        ]
    primerWriter.writerow(header)
    for clone in searchResults:
        print 'clonename=',clone
        #print searchResults[clone]
        length = searchResults[clone][2] - searchResults[clone][1]
        start = searchResults[clone][1]
        seq = primerRun(searchResults[clone][3])
        seq.format(start, length)
        seq.search(clone, primerWriter)
        os.remove('primer3_temp.txt')
    primerData.close()
    
    
if __name__ == '__main__':
    main()

